#!/usr/bin/python

import unittest
import os, rrdtool, datetime, time
import pyrrd

class PyRRDTest(unittest.TestCase):

    def setUp(self):
        self.testRRD = "test.rrd"
        self.testRRD2 = "test2.rrd"

    def testDS(self):
        self.assertEqual("DS:incoming:COUNTER:25:0:U",
                str(pyrrd.DS("incoming", pyrrd.COUNTER, 25, 0, None)))
        self.assertEqual("DS:incoming:COUNTER:25:U:U",
                str(pyrrd.DS("incoming", pyrrd.COUNTER, 25)))
        self.assertEqual("DS:incoming:COUNTER:20:U:U",
                str(pyrrd.DS("incoming", pyrrd.COUNTER, step=10)))
        self.assertEqual("DS:test:COMPUTE:blablabla",
                str(pyrrd.DS("test", pyrrd.COMPUTE, rpn="blablabla")))

    def testRRA(self):
        self.assertEqual("RRA:AVERAGE:0.1:1:8640",
                str(pyrrd.RRA(pyrrd.AVERAGE, 0.1, 1, 8640)))
        self.assertEqual("RRA:AVERAGE:0.5:1:8640",
                str(pyrrd.RRA(pyrrd.AVERAGE, None, 1, 8640)))
        self.assertEqual("RRA:MIN:0.5:1:720",
                str(pyrrd.RRA(pyrrd.MIN, step=120, interval=pyrrd.DAY)))
        self.assertEqual("RRA:MIN:0.5:5:1008",
                str(pyrrd.RRA(pyrrd.MIN, step=120, interval=pyrrd.WEEK, maxRows=1200)))
        self.assertEqual("RRA:MIN:0.5:15:1488",
                str(pyrrd.RRA(pyrrd.MIN, step=120, interval=pyrrd.MONTH)))
        #self.assertEqual("RRA:MIN:0.5:180:1464",
        #        str(pyrrd.RRA(pyrrd.MIN, step=120, interval=pyrrd.YEAR)))
        # why "RRA:MIN:0.5:176:1497" ?

    def testTimestamp(self):
        now = pyrrd.Timestamp.now()
        nowStamp = str(int(time.mktime(now.timetuple())))
        self.assertEqual(nowStamp, str(now))

    def createTest1(self):
        try:
            os.remove(self.testRRD)
        except os.error:
            pass
        rrdtool.create(self.testRRD,
                "--start", "now-1year",
                "--step", "10",
                "DS:incoming:COUNTER:25:0:U",
                "DS:outgoing:COUNTER:25:0:U",
                "DS:test:GAUGE:25:0:U",
                "RRA:AVERAGE:0.1:1:8640",
                "RRA:AVERAGE:0.1:50:1210",
                "RRA:AVERAGE:0.1:223:1202",
                "RRA:AVERAGE:0.1:2635:1201",
                "RRA:MIN:0.1:1:8640",
                "RRA:MIN:0.1:50:1210",
                "RRA:MIN:0.1:223:1202",
                "RRA:MIN:0.1:2635:1201",
                "RRA:MAX:0.1:1:8640",
                "RRA:MAX:0.1:50:1210",
                "RRA:MAX:0.1:223:1202",
                "RRA:MAX:0.1:2635:1201",
                "RRA:LAST:0.1:1:8640",
                "RRA:LAST:0.1:50:1210",
                "RRA:LAST:0.1:223:1202",
                "RRA:LAST:0.1:2635:1201")

    def updateTest1(self):
        now = pyrrd.Timestamp.now()
        delta = datetime.timedelta(seconds=10)
        time = now - datetime.timedelta(days=32)
        for i in xrange(0, 8640):
            rrdtool.update(self.testRRD, str(time))
        rrdtool.update(self.testRRD,
                "now-1month@2635000:263500:10000",
                "now-1week@223000:22300:1000",
                "now-1day@50000:5000:100",
                "now-1hour@1000:100:10")

    def createTest2(self):
        try:
            os.remove(self.testRRD2)
        except os.error:
            pass
        return pyrrd.create(filename=self.testRRD2,
                start="now-1year",
                step=10,
                ds = [
                pyrrd.DS("incoming", pyrrd.COUNTER, 25, 0, None),
                pyrrd.DS("outgoing", pyrrd.COUNTER, 25, 0, None),
                pyrrd.DS("test", pyrrd.GAUGE, 25, 0, None)],
                rra = [
                pyrrd.RRA(pyrrd.AVERAGE, 0.1, 1, 8640),
                pyrrd.RRA(pyrrd.AVERAGE, 0.1, 50, 1210),
                pyrrd.RRA(pyrrd.AVERAGE, 0.1, 223, 1202),
                pyrrd.RRA(pyrrd.AVERAGE, 0.1, 2635, 1201),
                pyrrd.RRA(pyrrd.MIN, 0.1, 1, 8640),
                pyrrd.RRA(pyrrd.MIN, 0.1, 50, 1210),
                pyrrd.RRA(pyrrd.MIN, 0.1, 223, 1202),
                pyrrd.RRA(pyrrd.MIN, 0.1, 2635, 1201),
                pyrrd.RRA(pyrrd.MAX, 0.1, 1, 8640),
                pyrrd.RRA(pyrrd.MAX, 0.1, 50, 1210),
                pyrrd.RRA(pyrrd.MAX, 0.1, 223, 1202),
                pyrrd.RRA(pyrrd.MAX, 0.1, 2635, 1201),
                pyrrd.RRA(pyrrd.LAST, 0.1, 1, 8640),
                pyrrd.RRA(pyrrd.LAST, 0.1, 50, 1210),
                pyrrd.RRA(pyrrd.LAST, 0.1, 223, 1202),
                pyrrd.RRA(pyrrd.LAST, 0.1, 2635, 1201)])

    def updateTest2(self):
        rrd = pyrrd.open(self.testRRD2)
        rrd.update(2635000, 263500, 10000, timestamp="now-1month")
        rrd.update(223000, 22300, 1000, timestamp="now-1week")
        rrd.update(50000, 5000, 100, timestamp="now-1day")
        rrd.update(1000, 100, 10, timestamp="now-1hour")

    def testCreate(self):
        self.createTest1()
        self.createTest2()
        info1 = rrdtool.info(self.testRRD)
        info2 = rrdtool.info(self.testRRD2)
        self.assertEqual(info1['ds'], info2['ds'])
        self.assertEqual(info1['rra'], info2['rra'])
        self.assertEqual(info1['rrd_version'], info2['rrd_version'])
        self.assertEqual(info1['step'], info2['step'])
        self.assertTrue(info2['last_update'] - info1['last_update'] < 5)
        #suppose that creation time differs not more than 5 seconds

    def testInit(self):
        rrd = pyrrd.RRD(self.testRRD)
        self.assertEqual(self.testRRD, rrd.filename)

    def testOpen(self):
        rrd = pyrrd.open(self.testRRD)
        self.assertEqual(self.testRRD, rrd.filename)

    def testParseUpdateParams(self):
        rrd = pyrrd.RRD(self.testRRD)
        self.assertEqual(('N:100:200', ''),
                rrd._parseUpdateParams((100, 200), {}))
        self.assertEqual(('t@100:200', ''),
                rrd._parseUpdateParams((100, 200), {'timestamp': 't'}))
        self.assertEqual(('1:100:200', ''),
                rrd._parseUpdateParams((100, 200), {'timestamp': '1'}))
        self.assertEqual(('N:100:200', 'a:b'),
                rrd._parseUpdateParams((), {'a': 100, 'b': 200}))
        self.assertEqual(('t@100:200', 'a:b'),
                rrd._parseUpdateParams((), {'a': 100, 'b': 200, 'timestamp': 't'}))
        self.assertEqual(('1:100:200', 'a:b'),
                rrd._parseUpdateParams((), {'a': 100, 'b': 200, 'timestamp': '1'}))

    def testUpdate(self):
        self.createTest1()
        self.createTest2()
        self.updateTest1()
        self.updateTest2()
        fetch1 = rrdtool.fetch(self.testRRD, 'LAST',
                '-s', 'now-1hour-1min', '-e', 'now-1hour+1min')
        fetch2 = rrdtool.fetch(self.testRRD2, 'LAST',
                '-s', 'now-1hour-1min', '-e', 'now-1hour+1min')
        self.assertEqual(fetch1, fetch2)

def test_suite():
    return unittest.makeSuite(PyRRDTest)

if __name__ == "__main__":
    unittest.TextTestRunner().run(test_suite())
